from abc import ABC
from typing import Dict, Optional, Any

from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
from langchain_core.callbacks import CallbackManagerForChainRun

from llmchat.chain.BaseChain import BaseChain


class ChatCompletionsChain(BaseChain, ABC):
    """TODO"""
    output_key: str = "answer"
    combine_docs_chain: BaseCombineDocumentsChain

    def _call(
            self,
            inputs: Dict[str, Any],
            run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        output: Dict[str, Any] = {}
        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
        new_inputs = inputs.copy()
        question = inputs["question"]
        answer = self.combine_docs_chain.invoke(inputs=new_inputs, config="")
        output[self.output_key] = answer

        return output


